import asyncio
import os
from openai import OpenAI
from dotenv import load_dotenv
from contextlib import AsyncExitStack

# 加载env文件，确保配置正确
load_dotenv()

class MCPClient:
    def __init__(self):
        """初始化MCP客户端"""
        self.exit_stack = AsyncExitStack()
        self.base_url=os.getenv("BASE_URL")
        self.model = os.getenv("MODEL")
        self.openai_api_key = os.getenv("OPENAI_API_KEY")

        if not self.openai_api_key:
            raise ValueError("OPENAI_API_KEY未设置")

        self.client=OpenAI(api_key=self.openai_api_key, base_url=self.base_url)

    async def connect_to_mock_server(self):
        """连接MCP服务器的连接"""
        print("MCP客户端初已始化，但未连接到服务器")

    async def process_query(self, query:str)->str:
        """调用OpenAI API处理用户查询"""
        message=[{"role": "system",
                  "content": "你是一个智能助手，默认使用中文回答问题."},
                 {"role": "user",
                  "content": query,}]
        try:
            response=await asyncio.get_event_loop().run_in_executor(
                            None,
                            lambda: self.client.chat.completions.create(model=self.model,messages=message)
                            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"调用OpenAI API时出错: {e}")


    async def chat_loop(self):
        """运行交互式聊天循环"""
        print("\nMCP客户端已启动！输入'quit'退出")
        while True:
            try:
                query=input("Query:").strip()
                if query.lower() == "quit":
                    break
                response=await self.process_query(query)
                print(f"\nResponse: {response}\n")
            except Exception as e:
                print(f"Error: {e}")
    async def cleanup(self):
        """清理资源"""
        await self.exit_stack.aclose()

async def main():
    client = MCPClient()
    try:
        await client.chat_loop()
    finally:
        await client.cleanup()

if __name__ == "__main__":
    asyncio.run(main())